#!/opt/ActivePython/bin/python
# daily.py
#--------------------------------------------------------------------------------------------
# Run this script from the same directory where treasury.db is located
#--------------------------------------------------------------------------------------------

__version__ = "$Revision: 1.8 $"
__author__ = 'Ramesh Balasubrmanian <ramesh@finpy.org>'
__doc__ = """
    Runs on weekdays at 6:00 pm US EST to check if new data is available in treasury website
    Any new data is loaded into treasury.db
    Usage: daily.py [/path/to/treasury.db]
    if path is not supplied treasury.db is assumed to be in the current directory. 
"""
import urllib
from xml.dom import minidom
import datetime
import time
import sys
from pysqlite2 import dbapi2 as sqlite

tagNames = ('BC_1MONTH', \
            'BC_3MONTH', \
            'BC_6MONTH', \
            'BC_1YEAR', \
            'BC_2YEAR', \
            'BC_3YEAR', \
            'BC_5YEAR', \
            'BC_7YEAR', \
            'BC_10YEAR', \
            'BC_20YEAR', \
            'BC_30YEAR')
url = 'http://www.ustreas.gov/offices/domestic-finance/debt-management/interest-rate/yield.xml'

def _setYieldData(ydata, yieldData, key):
    for y in ydata:
        for ycn in y.childNodes:
            if ycn.firstChild != None and ycn.nodeName in tagNames:
                try:
                    yieldData[key][list(tagNames).index(ycn.nodeName)] = \
                        '%.2f' % float(ycn.firstChild.nodeValue)
                except ValueError:
                    yieldData[key][list(tagNames).index(ycn.nodeName)] = ''
                                
def _fetchCurrent(url, yieldData):
    xmldata = urllib.urlopen(url).read()
    doc = minidom.parseString(xmldata)
    dates = doc.childNodes[0].getElementsByTagName('G_NEW_DATE')    
    for newDate in dates:
        key = newDate.getElementsByTagName('BID_CURVE_DATE')[0].firstChild.nodeValue
        keyDate = datetime.date(*time.strptime(key.lower(), "%d-%b-%y")[0:3])
        keyDateOrdinal = repr(keyDate.toordinal())
        yieldData[keyDateOrdinal] = ['', '', '', '', '', '', '', '', '', '', '']
        ydata = newDate.getElementsByTagName('G_BC_CAT')
        _setYieldData(ydata, yieldData, keyDateOrdinal)
    
def checkData(dbf = 'treasury.db'):
    connection = sqlite.connect(dbf)
    cursor = connection.cursor()
    sql = """select max(record_date) 
             from yield_curve""" 
    cursor.execute(sql)
    row = cursor.fetchone()
    if row == None:
        print 'Yield Curve data not available in db. Unexpected error'
    else:
        print 'Most recent record_date in database is', datetime.date.fromordinal(row[0])
        return row[0]

def getData(maxRecordDate = 0):
    """
    fetch current month yield curve data from treasury web site.
    """
    yieldData = {}
    _fetchCurrent(url, yieldData)
    for k, v in yieldData.items():
        if v == ['', '', '', '', '', '', '', '', '', '', ''] or int(k) <= int(maxRecordDate):
            yieldData.pop(k,v)
    return yieldData

def writeToDb(yieldData, dbf = 'treasury.db'):
    connection = sqlite.connect(dbf)
    cursor = connection.cursor()
    sqls = ["insert into yield_curve values(" +  k + ',' + ','.join([rate for rate in v ]) + ")" \
            for k, v in yieldData.items()] 
    count = [cursor.execute(sql) for sql in sqls]
    print len(count),' records added to treasury.db'
    connection.commit()
    cursor.close()
    connection.close()

if __name__ == '__main__':
	print __doc__ 
	if len(sys.argv) == 2:
		writeToDb(getData(checkData(sys.argv[1])), sys.argv[1])
	else:
		writeToDb(getData(checkData()))
